#pragma once

#include "MNN/Interpreter.hpp"
#include <string>
#include <array>

class Model {
public:
    /**
     * @brief Constructs a new Model object from a given model path.
     *
     * @param modelPath The file path to the model file.
     */
    Model(const std::string& modelPath);

    /**
     * @brief Destructor for the Model class.
     */
    ~Model();

    /**
    * @brief Fills the input tensor with the given data.
    *
    * This method fills the input tensor of the model with the provided float array. The array is expected to have a size
    * of 128 elements, each containing an array of 6 floats.
    *
    * @param inputs A pointer to a float array containing the data to be filled into the input tensor.
    */
    void fillInputTensor(const float inputs[128][6]);
    
    /**
     * @brief Runs inference and stores the output in the provided array.
     *
     * This method performs inference on the current model configuration using the data
     * filled into the input tensor. The result is stored in the provided output array.
     *
     * @note This method should be called after method \ref fillInputTensor.
     *
     * @param output A reference to an array of size 6 where the output will be stored.
     */
    void predict(std::array<float, 6>& output);

    /**
     * @brief Retrieves the label associated with a given index.
     *
     * This method returns the label string corresponding to the provided index.
     *
     * @param idx The index of the label to retrieve.
     * @return const std::string& A constant reference to the label string.
     */
    inline const std::string& getLabel(int idx) const {
        return labels[idx];
    }

    /**
     * @brief Returns the index of the maximum value in the array.
     *
     * This method finds and returns the index of the largest float value
     * within the given 6-element array.
     *
     * @param arr The input array containing 6 float values.
     * @return size_t The index of the maximum value in the array.
     */
    static inline size_t argmax(const std::array<float, 6>& arr) {
        return std::distance(arr.begin(), std::max_element(arr.begin(), arr.end()));
    }

private:
    MNN::Interpreter* interpreter; /*!< A pointer to the MNN interpreter. */
    MNN::Session* session; /*!< A pointer to the MNN session. */
    MNN::Tensor* inputTensor; /*!< A pointer to the input tensor. */
    MNN::Tensor* outputTensor; /*!< A pointer to the output tensor. */
    static const std::string labels[6]; /*!< An array of string labels for classification results. */
};
